import cPickle

import matplotlib
from matplotlib import pyplot
import minuit
import numpy
from scipy.integrate import quad

import fitter._minuit.tex
import fitter._minuit.signals

class GaussPlusPolynomialFit(object):
    max_polyorder = 10

    rcprops = {'axes.formatter.limits' : (-3, 3),
               'font.size': 24.0,
               'text.usetex': True,
               'text.latex.preamble': [r'\usepackage{amsmath}'],
               'xtick.labelsize': 'large',
               'xtick.direction': 'out',
               'ytick.labelsize': 'large',
               'ytick.direction': 'out',
               'figure.figsize': (12.0, 6.0),
               }

    figprops = {'axes': (0.13, 0.1425, 0.5, 0.75),
                'table_x': 1.69,
                'table_y': 0.95,
                }

    tex_labels = {'A' : r'A',
                  'mu': r'\mu',
                  'sigma': r'\sigma',
                  'p0': r'p_0',
                  'p1': r'p_1',
                  'p2': r'p_2',
                  'p3': r'p_3',
                  'p4': r'p_4',
                  'p5': r'p_5',
                  'p6': r'p_6',
                  'p7': r'p_7',
                  'p8': r'p_8',
                  'p9': r'p_9',
                }

    tex_header=r'$A\exp\left[\frac{(x - \mu)^2}{\sigma^2}\right]$'
    signal_fields = "A mu sigma".split()

    def __init__(self, hist, polyorder=3, loudness=0):
        self.m = minuit.Minuit(self._chi2)

        for key in self.m.values:
            self.m.values[key] = 1

        self.polyorder = polyorder
        self.loudness = loudness
        self.m.printMode = loudness

        self.hist = hist
        
        self.x = numpy.array([hist.bin_center(i) for i in xrange(hist.nbins)])
        self.y = hist.bins()[:]
        self.yunc = numpy.sqrt(self.y)
        self.yunc = numpy.where(self.yunc > 0, self.yunc, numpy.ones(len(self.yunc)))

        self.fit_x = self.x
        self.fit_y = self.y
        self.fit_yunc = self.yunc
        
        self.plot_bgonly = False

        for i in xrange(self.polyorder+1, GaussPlusPolynomialFit.max_polyorder):
            key = 'p{0}'.format(i)
            self.m.values[key] = 0.0
            self.m.fixed[key] = True
            
            
        self.ncalls = 0
        self._og_fixed = {}

        self.initial_chi2 = None
        self.final_chi2 = None
        

    def _signal(self, x, A, mu, sigma):
        return fitter._minuit.signals.gauss(x, A, mu, sigma)

    def _bg(self, x, p0, p1, p2, p3, p4, p5, p6, p7, p8, p9):
        return numpy.polynomial.chebyshev.chebval(
            x, [p0, p1, p2, p3, p4, p5, p6, p7, p8, p9])

    def _model(self, x, A, mu, sigma, p0, p1, p2, p3, p4, p5, p6, p7, p8, p9):
        return self._signal(x, A, mu, sigma) \
            + self._bg(x, p0, p1, p2, p3, p4, p5, p6, p7, p8, p9)

    def _chi2(self, A, mu, sigma, p0, p1, p2, p3, p4, p5, p6, p7, p8, p9):
        chi = (self._model(self.fit_x, A, mu, sigma, 
                           p0, p1, p2, p3, p4, p5, p6, p7, p8, p9) 
               - self.fit_y) / self.fit_yunc
        
        chisq = numpy.sum(chi**2)

        if self.loudness > 0 and self.ncalls % 100 == 0:
            print self.ncalls, chisq
            
        return chisq

    def signal(self):
        return self._signal(self.fit_x,
                            self.m.values['A'], 
                            self.m.values['mu'],
                            self.m.values['sigma'])
    def bg(self):
        return self._bg(self.fit_x,
                        self.m.values['p0'],
                        self.m.values['p1'],
                        self.m.values['p2'],
                        self.m.values['p3'],
                        self.m.values['p4'],
                        self.m.values['p5'],
                        self.m.values['p6'],
                        self.m.values['p7'],
                        self.m.values['p8'],
                        self.m.values['p9'],
                        )

    def model(self):
        return self._model(self.fit_x, **self.m.values)

    def chi2(self):
        return self._chi2(**self.m.values)
    
    def fit(self, low, high, **kwargs):
        """
        Fits the internally-stored histogram over the range (low, high).

        Arguments
        ---------
        low  : float
            The lower bound for the fit
        high : float
            The upper bound for the fit

        
        Keyword Arguments
        -----------------        
        fixed     : sequence of parameter names
            Fix the supplied parameters to their current values during the next fit
        signal_range : sequence of float of length 2
            The range over which the signal exists; if this argument is supplied, 
            ``signal_range`` will be omitted and the remaining bins will be fit to 
            the background function; after the fit is successful, the background 
            will be fixed prior to the signal fit.

        """


        self.fit_range = (low, high)

        lower_limit = self.hist.bin_index(low)
        upper_limit = self.hist.bin_index(high)
 
        if upper_limit < lower_limit:
            msg = 'Upper limit ({0} is less than lower limit ({1})'
            raise ValueError(msg.format(high, low))

        for par_name in self.m.values:
            if par_name in kwargs:
                self.m.values[par_name] = kwargs[par_name]
                del kwargs[par_name]

        self.initial_values = {}
        self.initial_values.update(self.m.values)

        if 'fixed' in kwargs:
            self._og_fixed.update(self.m.fixed)
            for par_name in kwargs['fixed']:
                self.m.fixed[par_name] = True
        
        if 'signal_range' in kwargs:
            signal_lower_limit = self.hist.bin_index(kwargs['signal_range'][0])
            signal_upper_limit = self.hist.bin_index(kwargs['signal_range'][1])
            
            self.release_background()

            if signal_lower_limit < lower_limit:
                msg = 'Signal lower limit ({0} is less than overall fit lower limit ({1})'
                raise ValueError(msg.format(kwargs['signal_range'][0], low))

            if signal_upper_limit > upper_limit:
                msg = 'Signal upper limit ({0} is greater than overall fit upper limit ({1})'
                raise ValueError(msg.format(kwargs['signal_range'][1], high))

            if signal_upper_limit < signal_lower_limit:
                msg = 'Signal upper limit ({0} is less than signal lower limit ({1})'
                raise ValueError(msg.format(kwargs['signal_range'][1], kwargs['signal_range'][0]))


            self.fit_x = numpy.concatenate((self.x[lower_limit:signal_lower_limit],
                                            self.x[signal_upper_limit:upper_limit])
                                           )
            self.fit_y = numpy.concatenate((self.y[lower_limit:signal_lower_limit],
                                            self.y[signal_upper_limit:upper_limit])
                                           )
        
            self.fit_yunc = numpy.concatenate((self.yunc[lower_limit:signal_lower_limit],
                                               self.yunc[signal_upper_limit:upper_limit])
                                              )    

            bginit_vals = {}
            bginit_vals.update(self.m.values)
            
            bginit_fixed = {}
            bginit_fixed.update(self.m.fixed)

            self.zero_signal()
            self.fix_signal()
 
            try:
                self.m.migrad()
            except minuit.MinuitError, e:
                msg = "Error encountered during background-only fit: {0}".format(str(e))
                raise minuit.MinuitError(msg)

            if self.plot_bgonly:
                print "Saving fit result to background-only fit to `{0}'...".format(bg_only_output)

                bg_only_output = 'bg_only_fit.png'
                self.plot(name=bg_only_output, curve_range=self.fit_range)
            

#            self.m.values['A']     = bginit_vals['A']
#            self.m.values['mu']    = bginit_vals['mu']
#            self.m.values['sigma'] = bginit_vals['sigma']

            self.bg_fit_x = self.fit_x
            self.bg_fit_y = self.fit_y
            self.bg_fit_yunc = self.fit_yunc
            
            self.m.fixed.update(bginit_fixed)
            
            self.fix_background()
 
            return
                
        self.fit_x = self.x[lower_limit:upper_limit]
        self.fit_y = self.y[lower_limit:upper_limit]
        self.fit_yunc = self.yunc[lower_limit:upper_limit]
        
        try:
            rt = self.m.migrad()
        except minuit.MinuitError, e:
            msg = "Error encountered during main fit: {0}".format(str(e))
            raise minuit.MinuitError(msg)

        self.initial_chi2 = self._chi2(**self.initial_values)
        self.final_chi2   = self.chi2()

        msg = "Fit converged with code {0}; initial chi^2 was {1}, final chi^2 is {2}"
        print msg.format(rt, self.initial_chi2, self.final_chi2)
        

    def zero_signal(self):        
        self.m.values['A']     = 0
        self.m.values['mu']    = 1
        self.m.values['sigma'] = 1        

    def fix_background(self):
        for i in xrange(GaussPlusPolynomialFit.max_polyorder):
            key = 'p{0}'.format(i)
            self.m.fixed[key] = True

    def fix_signal(self):
        for f in GaussPlusPolynomialFit.signal_fields:
            self.m.fixed[f] = True

    def release_signal(self):
        for f in GaussPlusPolynomialFit.signal_fields:
            self.m.fixed[f] = False
        
    def release_background(self):
        for i in xrange(GaussPlusPolynomialFit.max_polyorder):
            key = 'p{0}'.format(i)
            self.m.fixed[key] = False

        for i in xrange(self.polyorder+2, GaussPlusPolynomialFit.max_polyorder):
            key = 'p{0}'.format(i)
            self.m.values[key] = 0.0
            self.m.fixed[key] = True

    def release_all(self):
        self.release_signal()
        self.release_background()

    def integrate(self):
        integral = quad(lambda x: self._model(x, **self.m.values))

        return integral[0] / self.hist.bin_width, integral[1]

    def integrate_signal(self):
        integral = quad(lambda x: self._signal(x, 
                                               self.m.values['A'], 
                                               self.m.values['mu'], 
                                               self.m.values['sigma']))

        return integral[0] / self.hist.bin_width, integral[1]

    def integrate_background(self):
        integral = quad(lambda x: self._signal(x, 
                                               self.m.values['A'], 
                                               self.m.values['mu'], 
                                               self.m.values['sigma']))

        return integral[0] / self.hist.bin_width, integral[1]        

    def plot(self, **kwargs):
        """
        Plots the result of the last fit performed.

        
        Keyword Arguments
        -----------------        
        curve_range : sequence of length 2 
            Plot the fitted curve over the range specified
        lw: int
            Set the linewidth of the fitted curve to ``lw``
        name        : string
            Use ``name`` as the output path for the plot image; extensions can be omitted
        title       : string
            Use ``title`` as the plot title.
            
        """

        default_mpl_rcparams = {}
        default_mpl_rcparams.update(matplotlib.rcParams)

        matplotlib.rcParams.update(self.rcprops)

        fig = pyplot.figure()

        axes = pyplot.axes(self.figprops['axes'])

        textbl = fitter._minuit.tex.FitParameterTeXTable()
        
        texgroup = [(self.tex_labels[p], 
                     self.m.values[p], self.m.errors[p]) 
                    for p in self.signal_fields]

        textbl.add_group(texgroup, self.tex_header)

        keys = ['p{0}'.format(i) for i in xrange(self.polyorder+1)]
        textbl.add_group([(self.tex_labels[key],
                          self.m.values[key], self.m.errors[key]) 
                          for key in keys], 
                         header=r'$T_{0}(x)$'.format(self.polyorder))
            

        if self.initial_chi2 is not None and self.final_chi2 is not None:
            textbl.add_chi2(self.initial_chi2, self.final_chi2)

        fptex = textbl.write()
                                
        pyplot.text(self.figprops['table_x'], 
                    self.figprops['table_y'], 
                    r"\input{" + fptex.name + "}",
                    fontsize=18, verticalalignment='top',
                    horizontalalignment='right',
                    transform=axes.transAxes,
                    bbox={'boxstyle':'square', 'ec':'k', 'color':'#FFFFFF'})



        if 'subtract_bg' in kwargs:
            y = self.fit_y - self._bg(self.fit_x, 
                                      self.m.values['p0'], 
                                      self.m.values['p1'], 
                                      self.m.values['p2'], 
                                      self.m.values['p3'], 
                                      self.m.values['p4'], 
                                      self.m.values['p5'], 
                                      self.m.values['p6'], 
                                      self.m.values['p7'], 
                                      self.m.values['p8'], 
                                      self.m.values['p9'])
        
            pyplot.errorbar(self.fit_x, y, yerr=self.fit_yunc, fmt='r.')

            x_curve = numpy.linspace(self.fit_range[0], self.fit_range[1], 10000)
            y_curve = self._signal(x_curve, *[self.m.values[f] for f in self.signal_fields])
            
        else:
            pyplot.errorbar(self.x, self.y, yerr=self.yunc, fmt='r.')
            
            if 'curve_range' not in kwargs:
                kwargs['curve_range'] = (min(self.fit_x), max(self.fit_x))

            x_curve = numpy.linspace(kwargs['curve_range'][0], kwargs['curve_range'][1],
                                     10000)

            y_curve = self._model(x_curve, **self.m.values)

        if 'lw' not in kwargs:
            kwargs['lw'] = 2

        pyplot.plot(x_curve, y_curve, 'b-', lw=kwargs['lw'])

        ylim = pyplot.ylim()
        pyplot.ylim(0, ylim[1])

        pyplot.xlim(self.hist.arange)

        if 'xlabel' in kwargs:
            pyplot.xlabel(kwargs['xlabel'])
        else:
            pyplot.xlabel(self.hist.xlabel)

        if 'ylabel' in kwargs:
            pyplot.ylabel(kwargs['ylabel'])
        else:
            pyplot.ylabel(self.hist._set_ylabel())

        if 'title' not in kwargs:
            kwargs['title'] = self.hist.title

        pyplot.title(kwargs['title'])

        if 'name' not in kwargs:
            kwargs['name'] = self.hist.name + '_fit'

        if kwargs['name'][-4:] not in ('.png', '.pdf'):
            path_out = kwargs['name'] + '.png'
            pyplot.savefig(path_out)
            path_out = kwargs['name'] + '.pdf'
            pyplot.savefig(path_out)
        else:
            pyplot.savefig(kwargs['name'])
            

